"""
    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    created by lane.chang on '24/10/2022'
    comment: 异步redis相关
"""
import asyncio
import aioredis

from config import Config


class AioRedis:
    """ 异步redis
    """
    def __init__(self):
        """
        """
        self.redis_host = Config.REDIS_OPTIONS.get('host', '127.0.0.1')
        self.redis_user = ''
        self.redis_password = ''
        self.connection_url = f"redis://{self.redis_user}:{self.redis_password}@{self.redis_host}/"
        self.redis_db = Config.REDIS_OPTIONS.get('db', 0)
        # 创建redis连接
        self.connection = None

    async def create_connection(self):
        """
        :return:
        """
        self.connection = aioredis.from_url(self.connection_url, db=Config.REDIS_OPTIONS.get('db', 0))

        return self.connection

    async def add_to_stream(self, data: dict, stream_channel):
        """ 将一条数据添加到队列
        :param data:
        :param stream_channel:
        :return:
        """
        try:
            if not stream_channel:
                return
            if not self.connection:
                self.connection = await self.create_connection()
            msg_id = await self.connection.xadd(name=stream_channel, id="*", fields=data)
            return msg_id

        except Exception as e:
            raise Exception(f"Error sending msg to stream => {e}")

    async def consume_stream(self, count: int, block: int, stream_channel):
        """读取队列中的消息，但是并不删除

        Args:
            count (int): _description_
            block (int): _description_
            stream_channel (_type_): _description_

        Returns:
            _type_: _description_
        """
        if not self.connection:
            self.connection = await self.create_connection()
        response = await self.connection.xread(streams={stream_channel: '0-0'}, count=count, block=block)

        return response

    async def delete_message(self, stream_channel, message_id):
        """成功消费数据后，调用此函数删除队列数据

        Args:
            stream_channel (_type_): _description_
            message_id (_type_): _description_
        """
        if not self.connection:
            self.connection = await self.create_connection()
        await self.connection.xdel(stream_channel, message_id)

    async def set(self, name, value):
        """ 设置键值
        :param name:
        :param value:
        :return:
        """
        if not self.connection:
            self.connection = await self.create_connection()
        return await self.connection.set(name, value)

    async def delete(self, name):
        """ 删除键值
        :param name:
        :return:
        """
        if not self.connection:
            self.connection = await self.create_connection()
        return await self.connection.delete(name)

    async def exists(self, name):
        """ 判断键值
        :param name:
        :return:
        """
        if not self.connection:
            self.connection = await self.create_connection()
        return await self.connection.exists(name)


async def main():
    """
    :return:
    """
    aio_redis = AioRedis()
    data = {'xiaoming5': 123}
    await aio_redis.add_to_stream(data=data, stream_channel='message_channel')

    # 从队列中拿出最新的1条数据
    data = await aio_redis.consume_stream(1, block=0, stream_channel='message_channel')
    print(data)

    # 从队列中拿出最新的1条数据
    data = await aio_redis.consume_stream(1, block=0, stream_channel='message_channel')
    print(data)

    # 轮询等待队列中的新消息
    response = await aio_redis.consume_stream(stream_channel="message_channel", count=1, block=0)
    if response:
        for stream, messagees in response:
            print('stream:', stream)
            for message in messagees:
                print('message: ', message)
                message_id = message[0]
                print('message_id: ', message_id)
                message_content = message[1]
                print('message_content: ', message_content)
                print('注意里面的键、值都变成了byte类型，需要进行解码:')
                message_content: dict
                print('message_content_decode: ', {k.decode('utf-8'): v.decode('utf-8') for k, v in message_content.items()})

                # 消费成功后删除队列中的消息
                await aio_redis.delete_message(
                    stream_channel='message_channel', message_id=message_id
                )

    await aio_redis.delete('message_channel')


if __name__ == '__main__':
    """
    """
    asyncio.run(main())



